#!/usr/bin/env python3

"""
Author: Yanrui Hu
Date: 2023/5/10
Description: 使用SSH功能模块进行脚本的传输和运行
Keywords: Python, SSH, paramiko, ssh.connect
Process:
- 读取配置文件中的信息(hostname, port, username, password)
- 使用文件中的信息传输脚本文件 autodeploy.py 至若干主机
- 使用文件中的信息在若干主机执行脚本文件, 从而实现若干主机的统一部署
References:
- https://docs.paramiko.org/en/stable/index.html
Usage:
- python3 -u ssh-autodeploy.py [options] COMMAND [args]
查看帮助：
- python3 -u ssh-autodeploy.py -h
- python3 -u ssh-autodeploy.py --help

初始化服务器上的日志仓库：
- python3 -u ssh-autodeploy.py init

在服务器上部署新的文件系统版本：
- python3 -u ssh-autodeploy.py deploy [-i | --image IMAGE] [-f | --force]

IMAGE不为空时, 在所有服务器上使用指定的 IMAGE 部署;
IMAGE为空时, 需要在servers.json文件中自行设置 image 键对应的值。

-f | --force 表示「强制部署」,
ATTENTION: "force-deploy" will not undo the current deploy and deploy the new image onto it.

未使用此选项时, 是自动部署模式, 在部署新的镜像时, 程序自动分析status.log文件,
计算期望部署情况和当前部署情况的差异(diff_dirs), 如果发现存在一些不必撤销的公共层,
则不会完全撤销上次镜像对应的文件系统的更改, 仅仅撤销必须撤销的一部分, 之后, 继续部署未部署的层,
直到目标镜像部署成功。

注意: 使用此选项可能会造成一些麻烦！
    推荐使用此选项的场景：想要结合两个不同的镜像的特点。这两个镜像对原始镜像的修改应当是「不相干的」
    比如, 现有 从原始镜像得到的镜像A和B, A与B对于原始镜像修改丝毫不同
    假设我们需要将A与B的修改结合起来, 得到一个新的镜像C
    我们可以在部署A(或者B) 之后, 使用-f选项部署B(或者A), 即可达到与部署C相同的效果
    这种方法可以无限次应用：A + B + C + D + ... = X
    A, B, C, D, ... 都是在一个原始镜像上的修改, 且这些修改之间基本上没有任何关联（否则会出现一些意想不到的效果）
    这有些类似于Git的Feature Branch的概念, 但是这里的「分支」是在文件系统层面上的, 而不是在代码层面上的

For example:
- python3 -u ssh-autodeploy.py deploy -i "docker.io/yanruihu/ubuntu:v1.1"
- python3 -u ssh-autodeploy.py deploy -i "docker.io/yanruihu/ubuntu:v1.1"

撤销服务器上的部署：
- python3 -u ssh-autodeploy.py undeploy

更多用法，请在命令行中使用 -h 或者 --help 选项查看帮助。
"""
import paramiko
from pprint import pprint
import sys, os, json, time
import argparse
import logging, subprocess
from paramiko.ssh_exception import NoValidConnectionsError


REPO_PATH = os.path.dirname(os.path.abspath(__file__))
FILE_SRC = f'{REPO_PATH}/autodeploy.py'
SERVERS_JSON = f"{REPO_PATH}/servers.json"
LOG_FILE = f'{REPO_PATH}/ssh-running.log'
# FILE_TARGET = '/home/user/autodeploy.py'

# ------------------------------------------------------------
# 配置日志
# ------------------------------------------------------------
# 创建一个 Logger 实例
running_logger = logging.getLogger('running')
running_logger.setLevel(logging.INFO)

# 创建一个 FileHandler 实例
file_handler = logging.FileHandler(f'{REPO_PATH}/ssh-running.log')
file_handler.setLevel(logging.INFO)

# 创建标准输出handler, 输出ERROR级别以上的日志
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.ERROR)


# 设置日志消息格式
file_formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d %(funcName)s - %(message)s'
)
file_handler.setFormatter(file_formatter)

console_formatter = logging.Formatter('%(levelname)s - %(message)s')
console_handler.setFormatter(console_formatter)

# 添加 handler 到 Logger 实例
running_logger.addHandler(file_handler)
running_logger.addHandler(console_handler)


# ------------------------------------------------------------
# 使用 SSH 执行命令
# ------------------------------------------------------------
ssh = paramiko.SSHClient()
policy = paramiko.AutoAddPolicy()
ssh.set_missing_host_key_policy(policy)


def ssh_exec_cmd(hostname, port, username, password, cmd):
    """使用SSH 在目标机器上执行命令"""
    try:
        ssh.connect(hostname=hostname, port=port, username=username, password=password)
        stdin, stdout, stderr = ssh.exec_command(cmd)
        exit_code = stdout.channel.recv_exit_status()

        if exit_code == 0:
            result = stdout.read().decode('utf-8')
            running_logger.info(
                f'ssh {username}@{hostname} -p {port} -t {cmd} success! stdout output:'
                + result
            )
        else:
            error = stderr.read().decode('utf-8')
            running_logger.error(
                'SSH exec_command error: '
                f'{username}@{hostname} -p {port} -t {cmd} \n' + error
            )
            exit(exit_code)

        stdin.close()
        time.sleep(1)
        # pprint(stdout.read().decode())
        ssh.close()
    except NoValidConnectionsError as e:
        s = f"SSH connection error: {str(e)}"
        running_logger.error(s)
        exit(1)
        # raise Exception(s) from e


def transfer_file(hostname, port, username, password, file_src, file_target):
    """
    通过 scp 向另外一台主机传输文件
     hostname: 主机名或者ip
     port: 端口号
     username: SSH登录主机的用户名
     password: 用户的登录口令
     file_src: 本机上的待传输文件
     file_target: 目标主机的存放地点（需要保证目录存在）
    """
    cmd = f'scp -pq -P {port} {file_src} {username}@{hostname}:{file_target}'
    res = subprocess.run(cmd, shell=True, capture_output=True)
    if res.returncode == 0:
        running_logger.info(
            f"scp success: {file_src} -> {username}@{hostname}:{file_target}"
        )
    else:
        running_logger.error(
            f"scp failed: {file_src} -> {username}@{hostname}:{file_target}"
        )
        running_logger.error(res.stderr.decode())
        exit(res.returncode)


def transfer_file_1(hostname, port, username, password, file_src, file_target):
    """
    通过 sftp 向另外一台主机传输文件
     hostname: 主机名或者ip
     port: 端口号
     username: SSH登录主机的用户名
     password: 用户的登录口令
     file_src: 本机上的待传输文件
     file_target: 目标主机的存放地点（需要保证目录存在）
    """
    ssh.connect(hostname=hostname, port=port, username=username, password=password)
    sftp = ssh.open_sftp()

    if not sftp:
        running_logger.critical("Object `sftp` is None")

    sftp.put(file_src, file_target)

    # 下载用get, 目前用不着
    # sftp.get(file_src, file_target)
    ssh.close()  # Close the connection.


def transfer_files(
    hostname, port, username, password, file_srcs: list, file_targets: list
):
    """通过 scp 向另外一台主机传输多个文件
    与 transfer_file 的函数签名基本一致,
     唯一的区别在于, 因为需要传输多个文件, 所以使用列表的形式"""

    for i, j in zip(file_srcs, file_targets):
        cmd = f"scp -pq -P {port} {i} {username}@{hostname}:{j}"
        res = subprocess.run(cmd, shell=True, capture_output=True)
        if res.returncode == 0:
            running_logger.info(f"scp success: {i} -> {username}@{hostname}:{j}")
        else:
            running_logger.error(f"scp failed: {i} -> {username}@{hostname}:{j}")
            running_logger.error(res.stderr.decode())
            exit(res.returncode)

    ssh.close()


def transfer_files_1(
    hostname, port, username, password, file_srcs: list, file_targets: list
):
    """通过 sftp 向另外一台主机传输多个文件
    与 transfer_file 的函数签名基本一致,
     唯一的区别在于, 因为需要传输多个文件, 所以使用列表的形式"""
    ssh.connect(hostname=hostname, port=port, username=username, password=password)
    sftp = ssh.open_sftp()

    if not sftp:
        running_logger.critical("Object `sftp` is None")

    for i, j in zip(file_srcs, file_targets):
        sftp.put(i, j)
        running_logger.info(f"sftp 传输文件 {i} 到 {j} 成功")

    ssh.close()


def transfer_dir(hostname, port, username, password, dir_src, dir_target):
    """通过 scp 向另外一台主机传输目录
    与 transfer_file 的函数签名基本一致,
     唯一的区别在于, 因为需要传输多个文件, 所以使用列表的形式"""
    cmd = f"scp -pqr -P {port} {dir_src} {username}@{hostname}:{dir_target}"
    res = subprocess.run(cmd, shell=True, capture_output=True)
    if res.returncode == 0:
        running_logger.info(
            f"scp success: {dir_src} -> {username}@{hostname}:{dir_target}"
        )
    else:
        running_logger.error(
            f"scp failed: {dir_src} -> {username}@{hostname}:{dir_target}"
        )
        running_logger.error(res.stderr.decode())
        exit(res.returncode)


def init_wrapper():
    """初始化服务器的包装器, 主要的作用是：
     - 传入 autodeploy.py 文件
     - 在目标servers上建立 .database 目录 (使用SSH执行 mkdir 命令）
     - 将初始化的三个log文件传入 每一个server的 .database 目录下
     - 运行 clear 命令，将日志文件中可能的旧数据清除掉
    PS: 以上的所有文件都存放在用户的家目录, root用户的家目录是 /root;
    其他用户的家目录在 /home下面, 名称与用户名相同"""
    # Step1: 读取 servers.json 文件,  获得每个服务器的登录信息, 赋值给 servers (List)
    with open(SERVERS_JSON, 'r') as f:
        servers = json.load(f)

    for item in servers:
        if item['username'] == 'root':
            HOME_PATH = "/root"
        else:
            HOME_PATH = f"/home/{item['username']}"
        # breakpoint()

        # 在家目录创建日志数据库目录 .database
        ssh_exec_cmd(
            hostname=item['hostname'],
            port=item['port'],
            username=item['username'],
            password=item.get('password', None),
            cmd=f"mkdir -p {HOME_PATH}/.database",
        )

        file_srcs = [
            f"{REPO_PATH}/.database/operations.log",
            f"{REPO_PATH}/.database/status.log",
            f"{REPO_PATH}/.database/last-oper.log",
            f'{REPO_PATH}/autodeploy.py',
            f'{REPO_PATH}/fileutils.py',
            f'{REPO_PATH}/log.py',
        ]
        file_targets = [
            f"{HOME_PATH}/.database/operations.log",
            f"{HOME_PATH}/.database/status.log",
            f"{HOME_PATH}/.database/last-oper.log",
            f"{HOME_PATH}/autodeploy.py",
            f"{HOME_PATH}/fileutils.py",
            f"{HOME_PATH}/log.py",
        ]

        transfer_files(
            hostname=item['hostname'],
            port=item['port'],
            username=item['username'],
            password=item.get('password', None),
            file_srcs=file_srcs,
            file_targets=file_targets,
        )

        # 执行 clear 命令
        ssh_exec_cmd(
            hostname=item['hostname'],
            port=item['port'],
            username=item['username'],
            password=item.get('password', None),
            cmd=f"python3 autodeploy.py clear --quiet",
        )

    return


def deploy_wrapper(args: argparse.Namespace):
    """执行deploy命令的包装器, 主要的作用是：
    - 读取 servers.json 文件
    - 在目标 servers 上执行 deploy 命令
    注意：部署的时候可能会带有 image选项 和 force 标志"""
    # Step1: 读取 servers.json 文件,  获得登录每个服务器需要的信息, 赋值给 servers (List)
    with open(SERVERS_JSON, 'r') as f:
        servers = json.load(f)

    # 需要判断 IMAGE 是否在命令行中指出, 如果 args.image 为 None, 则需要检查每一个 server 的 image 是否可访问
    # 如果args.image 不为 None, 则不需要检查每一个server的 image
    if args.image is None:
        for item in servers:
            if 'image' not in item.keys():
                running_logger.critical(
                    f"No image specified for server [{item['hostname']}] in servers.json"
                )
                print(
                    "请检查servers.json 文件, 确保每一个server都有image值.\n或者可以使用 `--image IMAGE`选项 为所有 server 指定同一镜像.",
                    file=sys.stderr,
                )
                sys.exit(1)

    if args.image:  # 统一镜像部署
        for item in servers:
            ssh_exec_cmd(
                hostname=item['hostname'],
                port=item['port'],
                username=item['username'],
                password=item.get('password', None),
                cmd=f"python3 autodeploy.py deploy --image={args.image} --target={args.target_dir} --quiet "
                + ('--force' if args.force else ''),  # 执行脚本, 镜像使用命令行参数
            )
    else:  # 个性化镜像部署
        for item in servers:
            ssh_exec_cmd(
                hostname=item['hostname'],
                port=item['port'],
                username=item['username'],
                password=item.get('password', None),
                cmd=f"python3 autodeploy.py deploy --image={item['image']} --target={args.target_dir} --quiet "
                + ('--force' if args.force else ''),  # 执行脚本, 镜像使用servers.json的自定义值
            )
    return


def undeploy_wrapper(args: argparse.Namespace):
    """执行 undeploy 命令的包装器, 主要的作用是：
    - 读取 servers.json文件
    - 在目标 server 上执行 undeploy 命令"""
    # Step1: 读取 servers.json 文件,  获得登录每个服务器需要的信息, 赋值给 servers (List)
    with open(SERVERS_JSON, 'r') as f:
        servers = json.load(f)

    for item in servers:
        ssh_exec_cmd(
            hostname=item['hostname'],
            port=item['port'],
            username=item['username'],
            password=item.get('password', None),
            cmd=f"python3 autodeploy.py undeploy",  # 执行 undeploy 命令
        )

    return


def clear_wrapper(args: argparse.Namespace):
    """执行 clear 命令的包装器, 主要的作用是：
    - 读取 servers.json文件
    - 在目标 server 上执行 clear 命令"""
    # Step1: 读取 servers.json 文件,  获得登录每个服务器需要的信息, 赋值给 servers (List)
    with open(SERVERS_JSON, 'r') as f:
        servers = json.load(f)

    for item in servers:
        ssh_exec_cmd(
            hostname=item['hostname'],
            port=item['port'],
            username=item['username'],
            password=item.get('password', None),
            cmd=f"python3 autodeploy.py clear --quiet",  # 执行 clear 命令
        )

    return


def purge_wrapper(args: argparse.Namespace):
    """执行 purge 命令的包装器, 主要的作用是：
    - 读取 servers.json文件
    - 在目标 server 上执行 purge 命令"""
    # Step1: 读取 servers.json 文件,  获得登录每个服务器需要的信息, 赋值给 servers (List)
    with open(SERVERS_JSON, 'r') as f:
        servers = json.load(f)

    for item in servers:
        ssh_exec_cmd(
            hostname=item['hostname'],
            port=item['port'],
            username=item['username'],
            password=item.get('password', None),
            cmd=f"python3 autodeploy.py purge",  # 执行 purge 命令
        )

    return


def init(args: argparse.Namespace):
    running_logger.info(f"Start init!")
    init_wrapper()  # 只能在第一次初始化时调用！
    return


def deploy(args: argparse.Namespace):
    running_logger.info(f'Start deploy!')
    deploy_wrapper(args)
    return


def undeploy(args: argparse.Namespace):
    running_logger.info(f'Start undeploy!')
    undeploy_wrapper(args)
    return


def clear(args: argparse.Namespace):
    running_logger.info(f'Start clear!')
    clear_wrapper(args)
    return


def purge(args: argparse.Namespace):
    running_logger.info(f'Start purge!')
    purge_wrapper(args)
    return


if __name__ == '__main__':
    desc = "使用SSH功能模块进行部署脚本的批量传输和批量运行"

    # 解析命令行参数
    # Create the top-level parser
    parser = argparse.ArgumentParser(description=desc, add_help=True)
    # parser.add_argument('--foo', action='store_true', help='foo help')

    subparsers = parser.add_subparsers(
        # help='help-可以使用的子命令',
        title="子命令列表",
        # description="desc-本软件的子命令介绍信息如下：",
        # prog="prog-用法信息",
    )

    # ------------------------------------------------------------
    # Add subparser for `init` sub-command.
    # ------------------------------------------------------------
    subparser_init = subparsers.add_parser(
        'init',
        help='初始化 servers.json 中各个服务器的日志仓库，上传部署过程依赖的脚本文件，仅仅需要执行一次！',
        description="""init 命令用于执行初始化操作,
        向服务器上传执行部署的脚本文件以及初始日志数据库, 为后续的部署和撤销部署做准备。
        这个命令应当只在初始化阶段执行一次, 后续不要执行！""",
    )

    subparser_init.set_defaults(func=init)

    # ------------------------------------------------------------
    # Add subparser for `deploy` sub-command.
    # ------------------------------------------------------------
    subparser_deploy = subparsers.add_parser(
        'deploy',
        help='使用 servers.json 中的信息部署镜像，此过程中每个服务器记录自己的部署日志。',
        description="""deploy 命令用于执行部署操作, 在服务器端执行部署脚本。
        部署分为两种模式：
        默认的自动模式以及使用 --force参数开启的强制模式。
        使用后一种模式需要先仔细阅读用法信息(在此脚本文件的开始部分)。否则后果自负！""",
    )

    subparser_deploy.add_argument(
        '--image',
        '-i',
        type=str,
        action='store',
        help='''如果未提供此选项, 部署时, 每个服务器使用的镜像将会从servers.json中读取。
        如果提供了此选项, 所有服务器上部署的镜像将由此选项的值统一指定。
         (If not provided, the images specified in servers.json
        for each server will be used to deploy. This argument
        is used to uniformly specify the image to be deployed in every server.)''',
    )
    subparser_deploy.add_argument(
        '--target-dir',
        '-t',
        type=str,
        action='store',
        help='部署操作的目标目录 (target dir), 默认为根目录 /',
        default="/",
    )
    subparser_deploy.add_argument(
        '--force',
        '-f',
        action='store_true',
        help='指定使用强制部署 (force deploy)模式, 默认为自动部署 (auto deploy)模式',
    )

    subparser_deploy.set_defaults(func=deploy)

    # ------------------------------------------------------------
    # Add subparser for `undeploy` sub-command.
    # ------------------------------------------------------------
    subparser_undeploy = subparsers.add_parser(
        'undeploy',
        help='撤销 servers.json 中所有服务器的部署，还原到部署前的状态。',
        description='''undeploy 命令用于撤销所有的部署, 没有任何选项。
        如需撤销单个server的部署。需要ssh进入单个server的命令行中进行操作。
        ''',
    )

    subparser_undeploy.set_defaults(func=undeploy)

    # ------------------------------------------------------------
    # Add subparser for `clear` sub-command.
    # ------------------------------------------------------------
    subparser_clear = subparsers.add_parser(
        'clear',
        help='清除 servers.json 中每一台服务器上的部署日志，数据一旦清除无法找回！部署结果的撤销依赖日志数据库，所以请谨慎操作！',
        description='clear 命令用于清除部署时产生的日志, 没有任何选项。',
    )
    subparser_clear.add_argument(
        '--quiet',
        '-q',
        action='store_true',
        help='不输出冗余的信息，不进行过多提问 (quiet)',
    )
    subparser_clear.set_defaults(func=clear)

    # ------------------------------------------------------------
    # Add subparser for `purge` sub-command.
    # ------------------------------------------------------------
    subparser_purge = subparsers.add_parser(
        'purge',
        help='清除 servers.json 中所有服务器上因为部署而拉取的所有镜像。(不会影响已部署的结果,设计此选项仅为节省空间)',
        description='purge 命令用于清除拉取的所有镜像, 以释放存储空间, 没有任何选项。此选项十分安全, 不会影响已部署的结果。',
    )
    subparser_purge.set_defaults(func=purge)

    # # 创建exec子命令
    # exec_parser = subparsers.add_parser(
    #     'exec',
    #     description="desc 测试",
    #     help="help-测试",
    # )
    # exec_parser.set_defaults(func=my_exec_function)

    if len(sys.argv) == 1:  # 如果没有提供任何参数, 则打印帮助信息
        parser.print_help()
        sys.exit()

    args = parser.parse_args()
    running_logger.info(args)
    print(f'{vars(args)=}')

    if 'func' in args:
        args.func(args)  # 执行子命令对应的函数
        exit()
